from pprint import pprint
import urllib
import zipfile
from collections import defaultdict

from task import *

# python2.6 -m examples.wordcount_task eval "top_words(10, 46, 1342, 27827)"

class GutembergDownload(Node):
    def __init__(self, idx):
        self.idx = idx
        filename = mkdir("gutemberg_data/{self.idx}.zip".format(**locals()))
        super(GutembergDownload, self).__init__(filename)
    def _run(self):
        # try latin1 at first
        url = "http://www.gutenberg.org/files/{self.idx}/{self.idx}-8.zip".format(**locals())
        _, headers = urllib.urlretrieve(url, self.name)
        if 'text/html' in headers['Content-Type']:
            url = url.replace('-8.', '.')
            _, headers = urllib.urlretrieve(url, self.name)
        if 'text/html' in headers['Content-Type']:
            raise IOError("Unable to retrieve '{0}'!".format(url))
    def _load(self):
        zf = zipfile.ZipFile(self.name)
        return zf.open(zf.namelist()[0])

class WordCount(Node):
    def __init__(self, etext):
        self.etext = etext
        filename = mkdir(etext.name.replace('.zip', '.wc'))
        super(WordCount, self).__init__(filename, [etext])
    def _run(self):
        count = defaultdict(int)
        for line in self.etext.load():
            for word in line.split():
                word = word.strip(',.-').lower()
                if word:
                    count[word] += 1
        pprint(dict(count), open(self.name, 'w+'))
    def _load(self):
        return eval(open(self.name).read())

class GroupCounts(Node):
    def __init__(self, counts):
        self.counts = counts
        filename = mkdir("gutemberg_data/"+key(sorted(counts))+".wc")
        super(GroupCounts, self).__init__(filename, counts)
    def _run(self):
        count = defaultdict(int)
        for icount in [c.load() for c in self.counts]:
            for word in icount:
                count[word] += icount[word]
        pprint(dict(count), open(self.name, 'w+'))
    def _load(self):
        return eval(open(self.name).read())


def top_words(num, *idxs):
    counts = [WordCount(GutembergDownload(idx)) for idx in idxs]
    group_count = GroupCounts(counts).load()
    for word, count in sorted(group_count.items(), key=lambda (w,c): (-c,w))[:num]:
        print count, word


if __name__ == "__main__" and "__file__" in globals():
    import sys
    args = sys.argv[1:]
    if "eval" in args[:1]:
        eval(" ".join(args[1:]))
    else:
        print "Running %r doctests..." % __file__
        import doctest
        doctest.testmod()
        print "DONE!"

